#include <linux/sched.h>
#include <linux/pgtable.h>
#include <linux/sched/task_stack.h>
#include <asm/switch_to.h>

static inline void finish_lock_switch(struct rq *rq)
{
    raw_spin_rq_unlock_irq(rq);
}

static inline void prepare_task(struct task_struct *next)
{
	/*
	 * Claim the task as running, we do this before switching to it
	 * such that any running task will have this set.
	 *
	 * See the smp_load_acquire(&p->on_cpu) case in ttwu() and
	 * its ordering comment.
	 */
	WRITE_ONCE(next->on_cpu, 1);
}

static inline void finish_task(struct task_struct *prev)
{
	/*
	 * This must be the very last reference to @prev from this CPU. After
	 * p->on_cpu is cleared, the task can be moved to a different CPU. We
	 * must ensure this doesn't happen until the switch is completely
	 * finished.
	 *
	 * In particular, the load of prev->state in finish_task_switch() must
	 * happen before this.
	 *
	 * Pairs with the smp_cond_load_acquire() in try_to_wake_up().
	 */
	smp_store_release(&prev->on_cpu, 0);
}

static inline void prepare_task_switch(struct rq *rq, struct task_struct *prev,
                                       struct task_struct *next)
{
    prepare_task(next);
}

static struct rq *finish_task_switch(struct task_struct *prev)
{
    struct rq *rq = task_rq(prev);
    unsigned int prev_state;

	/*
	 * The previous task will have left us with a preempt_count of 2
	 * because it left us after:
	 *
	 *	schedule()
	 *	  preempt_disable();			// 1
	 *	  __schedule()
	 *	    raw_spin_lock_irq(&rq->lock)	// 2
	 *
	 * Also, see FORK_PREEMPT_COUNT.
	 */
    prev_state = READ_ONCE(prev->__state);

    finish_task(prev);
    finish_lock_switch(rq);

    if (prev_state == TASK_DEAD)
    {
		task_dead(prev);

        put_task_finish_switch(prev);
    }

    return rq;
}

/*
 * context_switch - switch to the new MM and the new thread's register state.
 */
static inline struct rq *context_switch(struct rq *rq, struct task_struct *prev,
                                        struct task_struct *next, struct rq_flags *rf)
{
    prepare_task_switch(rq, prev, next);

    /*
     * For paravirt, this is coupled with an exit in switch_to to
     * combine the page table reload and the switch backend into
     * one hypercall.
     */
    arch_start_context_switch(prev);

    /* Here we just switch the register state and the stack. */
    switch_to(prev, next, prev);
    barrier();

    return finish_task_switch(prev);
}
